package org.study.ai.controller;

import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.model.Generation;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

/**
 * @author xl-9527
 * @since 2024/10/25
 **/
@RestController
@RequestMapping("ai/chat-models")
public class AiRequestController {

    private final OllamaChatModel ollamaChatModel;

    public AiRequestController(final OllamaChatModel ollamaChatModel) {
        this.ollamaChatModel = ollamaChatModel;
    }

    @GetMapping("ollama/dialogue")
    public Object dialogue(@RequestParam(name = "msg") String msg) {
        UserMessage userMessage = new UserMessage(msg);
        ChatResponse callResponse = ollamaChatModel.call(
                new Prompt(
                        userMessage,
                        OllamaOptions.builder()
                                .temperature(0.4)
                                .build()
                )
        );
        Generation result = callResponse.getResult();
        return result.getOutput().getContent();
    }
}
